import  torch
from    torch import optim, nn
# import  visdom
import  torchvision
from    torch.utils.data import DataLoader
from datam import *
from    resnet import ResNet18
from config import config,getData,getPolemon
import wandb
import yaml


data_dir = config.data_dir
batchsz = config.BATCH_SIZE
lr = config.lr
epochs = config.max_epochs
torch.manual_seed(1234)
train_loader,val_loader,test_loader = getPolemon()
# train_loader,val_loader,test_loader = getData()

def train(model, device, train_loader, optimizer, epoch,criteon):
    model.train()
    for step, (x, y) in enumerate(train_loader):
        x, y = x.to(device), y.to(device)
        model.train()

        logits = model(x)
        loss = criteon(logits, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        wandb.log({
            'epoch': epoch,
            'lr': lr,
            'loss': loss.item(),
        })

def evalute(model, device,loader):
    model.eval()
    correct = 0
    total = len(loader.dataset)
    for x,y in loader:
        x,y = x.to(device), y.to(device)
        with torch.no_grad():
            logits = model(x)
            pred = logits.argmax(dim=1)
        correct += torch.eq(pred, y).sum().float().item()
    return correct / total

def compare(model,val_acc,best_acc,epoch):
    if val_acc > best_acc:
        best_epoch = epoch
        best_acc = val_acc
        torch.save(model.state_dict(), 'best.mdl')
        wandb.log({
            'val_acc': val_acc,
        })
    return best_acc,best_epoch


def main():
    device = torch.device('cuda')
    model = ResNet18(5).to(device)
    optimizer = optim.Adam(model.parameters(), lr=lr)
    criteon = nn.CrossEntropyLoss()
    best_acc, best_epoch = 0, 0
    for epoch in range(epochs):
        train(model, device, train_loader, optimizer, epoch,criteon)
        val_acc = evalute(model, device,val_loader)
        best_acc,best_epoch = compare(model,val_acc,best_acc,epoch)

    print('best acc:', best_acc, 'best epoch:', best_epoch)
    model.load_state_dict(torch.load('best.mdl'))
    print('loaded from ckpt!')
    test_acc = evalute(model, device,test_loader)
    print('test acc:', test_acc)



if __name__ == '__main__':
    import time
    # wandb.init(project=config.project_name,name="getPolemon "+time.strftime("%Y_%m_%d_%H_%M_%S", time.localtime(time.time())))
    wandb.init(project=config.project_name,name="getData ")
    main()
    wandb.finish()